#ifndef RAIS_BSTREE_H
#define RAIS_BSTREE_H

#include <cstddef>
#include <cstdint>
#include <utility>

namespace rais {
namespace study {

namespace bstree_detail {
using std::size_t;
using std::uint8_t;
using std::forward;

template <typename T>
struct bsnode {
	T data;
	bsnode* left  = nullptr, 
	      * right = nullptr;
};

template <typename T, typename K>
constexpr K default_keyof(const T& data) { return K(data); }

template <typename T, typename K = T, K(*KeyOf)(const T&) = default_keyof<T, K>>
struct bstree {
	using data_t = T;
	using key_t = K;
	using node_t = bsnode<T>;
	using key_map_t = K(*)(const T&);
	static constexpr key_map_t keyof = KeyOf;

private:

	node_t* root = nullptr;
	size_t size = 0;

	node_t** find_node(key_t k) noexcept{
		node_t** pos = &root;
		while(*pos != nullptr) {
			key_t poskey = keyof((*pos)->data);
			if(k == poskey) break;
			else if(k < poskey) pos = &(*pos)->left;
			else                pos = &(*pos)->right;
		}
		return pos;
	}

	template <typename F>
	static void for_each_recursive(F&& f, node_t* pos) {
		if(pos == nullptr) return;
		for_each_recursive(f, pos->left);
		f(pos->data);
		for_each_recursive(f, pos->right);
	}

	template <typename F>
	static void for_each_recursive(F&& f, const node_t* pos) {
		if(pos == nullptr) return;
		for_each_recursive(f, pos->left);
		f(pos->data);
		for_each_recursive(f, pos->right);
	}

public:
	bstree() {}

	node_t* get_root() noexcept{return root;}
	const node_t* get_root() const noexcept{return root;}
	size_t get_size() const noexcept{return size;}

	data_t* at(key_t k) {
		node_t* pos = root;
		while(pos != nullptr) {
			key_t poskey = keyof(pos->data);
			if(k == poskey) &pos->data;
			else if(k < poskey) pos = pos->left;
			else                pos = pos->right;
		}
		return nullptr;
	}

	bool insert(const data_t& data) {
		if(root == nullptr) {
			root = new node_t{data};
			size++;
			return true;
		}
		node_t** pos = find_node(keyof(data));
		if(*pos != nullptr) return false;
		*pos = new node_t{data};
		size++;
		return true;
	} 

	bool insert(data_t&& data) {
		if(root == nullptr) {
			root = new node_t{forward<data_t>(data)};
			size++;
			return true;
		}
		node_t** pos = find_node(keyof(data));
		if(*pos != nullptr) return false;
		*pos = new node_t{forward<data_t>(data)};
		size++;
		return true;
	}

	bool erase(key_t k) {
		if(size == 0) return false;
		node_t** target_pos = find_node(k);
		if(*target_pos == nullptr) return false;
		//cases:
		// 0b00         : no child node   -> directly remove the node
		// 0b01 or 0b10 : one child node  -> let the child node replace the parent node  
		// 0b11         : two child nodes -> let the left node replace the parent node, and let the right node merge to the left branch's biggest child leaf
		uint8_t case_flag = (((*target_pos)->left != nullptr) << 1) | ((*target_pos)->right != nullptr);
		switch(case_flag) {
			case 0b00: {
				delete *target_pos;
				*target_pos = nullptr;
				break;
			}
			case 0b10: {
				node_t* child = (*target_pos)->left;
				delete *target_pos;
				*target_pos = child;
				break;
			}
			case 0b01: {
				node_t* child = (*target_pos)->right;
				delete *target_pos;
				*target_pos = child;
				break;
			}
			case 0b11: {
				node_t* left = (*target_pos)->left, 
				    * left_max = (*target_pos)->left;
				while(left_max->right != nullptr) {
					left_max = left_max->right;
				}
				left_max->right = (*target_pos)->right;
				delete *target_pos;
				*target_pos = left;
				break;
			}
		}
		size--;
		return true;
	}

	bool erase(const data_t& data) {
		return erase(keyof(data));
	}

	template <typename F>
	void for_each(F&& f) {
		for_each_recursive(f, root);
	}

	template <typename F>
	void for_each(F&& f) const{
		for_each_recursive(f, root);
	}

	
	bool contains(key_t key) const noexcept{
		node_t* pos = root;
		while(pos != nullptr) {
			key_t poskey = keyof(pos->data);
			if(key == poskey) return true;
			else if(key < poskey) pos = pos->left;
			else                  pos = pos->right;
		}
		return false;
	}
	bool contains(const data_t& data) const noexcept{
		return contains(keyof(data));
	}

	template <typename OStreamT>
	friend OStreamT& operator<<(OStreamT& os, bstree& self) {
		if(self.root == nullptr) return os << "[]";
		os << '[';
		for_each_recursive([&os](auto& data){os << data << ' ';}, self.root);
		return os << ']';
	}

	template <typename OStreamT>
	friend OStreamT& operator<<(OStreamT& os, const bstree& self) {
		if(self.root == nullptr) return os << "[]";
		os << '[';
		for_each_recursive([&os](auto& data){os << data << ' ';}, self.root);
		return os << ']';
	}

}; //struct bstree

} //namespace bstree_detail



} //namespace study
} //namespace rais

#endif //RAIS_BSTREE_H